{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "feb31cc6",
   "metadata": {},
   "source": [
    "# Running Agent as an Iterator\n",
    "\n",
    "It can be useful to run the agent as an iterator, to add human-in-the-loop checks as needed.\n",
    "\n",
    "To demonstrate the `AgentExecutorIterator` functionality, we will set up a problem where an Agent must:\n",
    "\n",
    "- Retrieve three prime numbers from a Tool\n",
    "- Multiply these together. \n",
    "\n",
    "In this simple problem we can demonstrate adding some logic to verify intermediate steps by checking whether their outputs are prime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8167db11",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import AgentType, initialize_agent\n",
    "from langchain.chains import LLMMathChain\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "from langchain_core.tools import Tool\n",
    "from langchain_openai import ChatOpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea6d45b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --upgrade --quiet  numexpr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7e41b9e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# need to use GPT-4 here as GPT-3.5 does not understand, however hard you insist, that\n",
    "# it should use the calculator to perform the final calculation\n",
    "llm = ChatOpenAI(temperature=0, model=\"gpt-4\")\n",
    "llm_math_chain = LLMMathChain.from_llm(llm=llm, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81e88aa5",
   "metadata": {},
   "source": [
    "Define tools which provide:\n",
    "\n",
    "- The `n`th prime number (using a small subset for this example)\n",
    "\n",
    "- The `LLMMathChain` to act as a calculator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "86f04b55",
   "metadata": {},
   "outputs": [],
   "source": [
    "primes = {998: 7901, 999: 7907, 1000: 7919}\n",
    "\n",
    "\n",
    "class CalculatorInput(BaseModel):\n",
    "    question: str = Field()\n",
    "\n",
    "\n",
    "class PrimeInput(BaseModel):\n",
    "    n: int = Field()\n",
    "\n",
    "\n",
    "def is_prime(n: int) -> bool:\n",
    "    if n <= 1 or (n % 2 == 0 and n > 2):\n",
    "        return False\n",
    "    for i in range(3, int(n**0.5) + 1, 2):\n",
    "        if n % i == 0:\n",
    "            return False\n",
    "    return True\n",
    "\n",
    "\n",
    "def get_prime(n: int, primes: dict = primes) -> str:\n",
    "    return str(primes.get(int(n)))\n",
    "\n",
    "\n",
    "async def aget_prime(n: int, primes: dict = primes) -> str:\n",
    "    return str(primes.get(int(n)))\n",
    "\n",
    "\n",
    "tools = [\n",
    "    Tool(\n",
    "        name=\"GetPrime\",\n",
    "        func=get_prime,\n",
    "        description=\"A tool that returns the `n`th prime number\",\n",
    "        args_schema=PrimeInput,\n",
    "        coroutine=aget_prime,\n",
    "    ),\n",
    "    Tool.from_function(\n",
    "        func=llm_math_chain.run,\n",
    "        name=\"Calculator\",\n",
    "        description=\"Useful for when you need to compute mathematical expressions\",\n",
    "        args_schema=CalculatorInput,\n",
    "        coroutine=llm_math_chain.arun,\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e660ee6",
   "metadata": {},
   "source": [
    "Construct the agent. We will use OpenAI Functions agent here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "21c775b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import hub\n",
    "\n",
    "# Get the prompt to use - you can modify this!\n",
    "# You can see the full prompt used at: https://smith.langchain.com/hub/hwchase17/openai-functions-agent\n",
    "prompt = hub.pull(\"hwchase17/openai-functions-agent\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ae7b104b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import create_openai_functions_agent\n",
    "\n",
    "agent = create_openai_functions_agent(llm, tools, prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "54e27bda",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import AgentExecutor\n",
    "\n",
    "agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a233fe4e",
   "metadata": {},
   "source": [
    "Run the iteration and perform a custom check on certain steps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "582d61f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `GetPrime` with `{'n': 998}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3m7901\u001b[0mChecking whether 7901 is prime...\n",
      "Should the agent continue (Y/n)?:\n",
      "y\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `GetPrime` with `{'n': 999}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3m7907\u001b[0mChecking whether 7907 is prime...\n",
      "Should the agent continue (Y/n)?:\n",
      "y\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `GetPrime` with `{'n': 1000}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3m7919\u001b[0mChecking whether 7919 is prime...\n",
      "Should the agent continue (Y/n)?:\n",
      "y\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `Calculator` with `{'question': '7901 * 7907 * 7919'}`\n",
      "\n",
      "\n",
      "\u001b[0m\n",
      "\n",
      "\u001b[1m> Entering new LLMMathChain chain...\u001b[0m\n",
      "7901 * 7907 * 7919\u001b[32;1m\u001b[1;3m```text\n",
      "7901 * 7907 * 7919\n",
      "```\n",
      "...numexpr.evaluate(\"7901 * 7907 * 7919\")...\n",
      "\u001b[0m\n",
      "Answer: \u001b[33;1m\u001b[1;3m494725326233\u001b[0m\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "\u001b[33;1m\u001b[1;3mAnswer: 494725326233\u001b[0mShould the agent continue (Y/n)?:\n",
      "y\n",
      "\u001b[32;1m\u001b[1;3mThe product of the 998th, 999th and 1000th prime numbers is 494,725,326,233.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "question = \"What is the product of the 998th, 999th and 1000th prime numbers?\"\n",
    "\n",
    "for step in agent_executor.iter({\"input\": question}):\n",
    "    if output := step.get(\"intermediate_step\"):\n",
    "        action, value = output[0]\n",
    "        if action.tool == \"GetPrime\":\n",
    "            print(f\"Checking whether {value} is prime...\")\n",
    "            assert is_prime(int(value))\n",
    "        # Ask user if they want to continue\n",
    "        _continue = input(\"Should the agent continue (Y/n)?:\\n\") or \"Y\"\n",
    "        if _continue.lower() != \"y\":\n",
    "            break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6934ff8e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
